import os
import json
import random
from sklearn.model_selection import train_test_split
from sklearn.model_selection import KFold

# 数据集根目录，需要根据实际情况修改
data_root = '../../dataset/brats2021'
# 输出的 JSON 文件路径
output_json_path = './jsons/brats2021_json_list.json'
# 定义折数，可根据需要修改
fold_num = 5

# 定义 BraTS 2021 数据集的标签信息
brats_labels = {
    "0": "Background",
    "1": "Non-enhancing tumor core",
    "2": "Peritumoral edema",
    "4": "GD-enhancing tumor"
}

# 定义模态信息
brats_modalities = {
    "0": "FLAIR",
    "1": "T1",
    "2": "T1ce",
    "3": "T2"
}

# 遍历数据集，收集图像和标签文件
image_files = []
label_files = []
for root, dirs, files in os.walk(data_root):
    for file in files:
        if file.endswith('_flair.nii.gz'):
            case_id = file.split('_flair.nii.gz')[0]
            image_paths = [
                os.path.join(root, f'{case_id}_{modality}.nii.gz')
                for modality in ['flair', 't1', 't1ce', 't2']
            ]
            label_path = os.path.join(root, f'{case_id}_seg.nii.gz')
            if all(os.path.exists(path) for path in image_paths + [label_path]):
                image_files.append(image_paths)
                label_files.append(label_path)

# 划分训练集和验证集
kf = KFold(n_splits=fold_num, shuffle=True, random_state=42)
train_val_data = []

# 构建训练集和验证集的数据列表
for fold, (_, val_index) in enumerate(kf.split(image_files)):
    for idx in val_index:
        train_val_data.append({
            "image": image_files[idx],
            "label": label_files[idx],
            "fold": fold
        })

# 构建 JSON 数据
json_data = {
    "description": "BraTS 2021 Brain Tumor Segmentation Dataset Split",
    "labels": brats_labels,
    "licence": "BraTS dataset license",
    "modality": brats_modalities,
    "name": "BraTS2021",
    "reference": "https://www.med.upenn.edu/cbica/brats2021.html",
    "release": "1.0",
    "tensorImageSize": "4D",
    "training": train_val_data
}

# 保存 JSON 文件
with open(output_json_path, 'w') as f:
    json.dump(json_data, f, indent=4)

print(f"JSON 文件已生成，保存路径为: {output_json_path}")